package org.lwj.logger.web;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

public class LogResponseWrapper extends HttpServletResponseWrapper {
	private ByteArrayOutputStream buffer = null;
	private ServletOutputStream out = null;
	private PrintWriter writer = null;

	public LogResponseWrapper(HttpServletResponse resp) throws IOException {
		super(resp);
		buffer = new ByteArrayOutputStream();
		out = new CachedServletOutputStream(buffer);
		writer = new PrintWriter(new OutputStreamWriter(buffer, this.getCharacterEncoding()));
	}

	@Override
	public ServletOutputStream getOutputStream() throws IOException {
		return out;
	}

	@Override
	public PrintWriter getWriter() throws UnsupportedEncodingException {
		return writer;
	}

	@Override
	public void flushBuffer() throws IOException {
		if (out != null) {
			out.flush();
		}
		if (writer != null) {
			writer.flush();
		}
	}

	@Override
	public void reset() {
		buffer.reset();
	}

	public String getCharacterEncoding() {
		String charsetName = super.getCharacterEncoding();
		charsetName = charsetName != null ? charsetName : "utf-8";
		return charsetName;
	}

	/**
	 * 获取响应内容
	 * 
	 * @return
	 * @throws Exception
	 */
	public String getContent() throws Exception {
		String charsetName = getCharacterEncoding();
		return getContent(charsetName);
	}

	/**
	 * 获取响应内容
	 * 
	 * @param charsetName
	 * @return
	 * @throws Exception
	 */
	public String getContent(String charsetName) throws Exception {
		return new String(getContentAsByteArray(), charsetName);
	}

	public byte[] getContentAsByteArray() throws IOException {
		flushBuffer();
		return buffer.toByteArray();
	}

	public class CachedServletOutputStream extends ServletOutputStream {
		private ByteArrayOutputStream baos = null;

		public CachedServletOutputStream(ByteArrayOutputStream stream) throws IOException {
			baos = stream;
		}

		@Override
		public void write(int b) throws IOException {
			baos.write(b);
		}

		@Override
		public void write(byte[] b) throws IOException {
			baos.write(b, 0, b.length);
		}

		public boolean isReady() {
			return false;
		}

	}
}